package rabbit.flt.plugins.springmvc;

import javax.servlet.ServletOutputStream;
import javax.servlet.http.HttpServletResponse;
import javax.servlet.http.HttpServletResponseWrapper;
import java.io.IOException;
import java.io.PrintWriter;
import java.util.function.Consumer;

public class ServletResponseWrapper extends HttpServletResponseWrapper {

    private OutputStreamWrapper streamWrapper;

    private PrintWriterWrapper writerWrapper;

    private Consumer<String> callBack;

    public ServletResponseWrapper(HttpServletResponse response, Consumer callBack) {
        super(response);
        this.callBack = callBack;
    }

    @Override
    public ServletOutputStream getOutputStream() throws IOException {
        if (null == streamWrapper) {
            streamWrapper = new OutputStreamWrapper(super.getOutputStream(), this);
        }
        return streamWrapper;
    }

    @Override
    public PrintWriter getWriter() throws IOException {
        if (null == writerWrapper) {
            writerWrapper = new PrintWriterWrapper(super.getWriter(), this);
        }
        return writerWrapper;
    }

    public Consumer<String> getCallBack() {
        return callBack;
    }
}
